# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com

import PricingEngineMod
from PricingEngineMod import *
reload(PricingEngineMod)

from math import pow, pi
from scipy.stats import norm
from scipy.integrate import quad as integ
from scipy import exp, log, sqrt, fft, r_
# Params: initial, type, rate, isPartial, isFFT, strike, tenor
#   HullWhite1Params: Lambda, Nu
#   HullWhite2Params: Lambda, Nu
#   HestonParams: Gamma, Kappa, V0, Mean
#   Corrs: Corr_Hest_HestVol, Corr_HW1_HW2, Corr_HW1_Hest, Corr_HW2_Hest, Corr_HW1_HestVol, Corr_HW2_HestVol
class HestonHullWhiteAnalytical(PricingEngine):

    def FourierFunction(k, alpha):
        val = 1 / (alpha ** 2 + alpha)
        return val
    
    def C1Integrand(self, RPARAM, type, omega):
        alpha   =   (-1) ** int(type)
        Bd = (exp(-self.HullWhite1Params_Lambda * RPARAM)-1.0) / self.HullWhite1Params_Lambda
        Bf = (exp(-self.HullWhite2Params_Lambda * RPARAM)-1.0) / self.HullWhite2Params_Lambda
    
        # Essentially D(w,r) is the same as D(w,tau), but now r is a variable 
        a           =   0.5 * self.HestonParams_Gamma ** 2
        b           =   (1j * omega + 2.0 - type) * self.Corr_Hest_HestVol * self.HestonParams_Gamma - self.HestonParams_Kappa
        c           =   -0.5 * (1j * alpha * omega + omega ** 2)
        d           =   sqrt(b ** 2 - 4.0 * a * c)
        r_plus      =   (-b + d)/(2.0 * a)
        r_minus     =   (-b - d)/(2.0 * a)
        g           =   r_minus / r_plus

        func_D_r    =   r_minus * (1.0 - exp(-d * RPARAM)) / (1 - g * exp(-d * RPARAM))

        # Now we compute Psi(r), see equation [18]. Psi(r) contains non-affine term: sqrt(v), we proxy it by phi_proxy_r, see equation [15] 
        h = (4.0 * self.HestonParams_Kappa * self.HestonParams_Mean) / self.HestonParams_Gamma ** 2;
              
        q_1 = self.HestonParams_Gamma**2 * (1.0 - exp(-self.HestonParams_Kappa)) / (4.0 * self.HestonParams_Kappa);
        l_1 = (4.0 * self.HestonParams_Kappa * self.HestonParams_V0 * exp(-self.HestonParams_Kappa)) / (self.HestonParams_Gamma**2 * (1.0 - exp(-self.HestonParams_Kappa)));
        h_1 = h + l_1;
        phi_1 = sqrt(q_1 * (l_1 - 1.0)+ q_1 * h + 0.5 * (q_1 * h) / h_1);
        
        a_a = sqrt(self.HestonParams_Mean - self.HestonParams_Gamma**2 / (8.0 * self.HestonParams_Kappa));
        b_b = sqrt(self.HestonParams_V0) - a_a;
        c_c = -log((phi_1 - a_a) / b_b);
        phi_proxy_r = a_a + b_b * exp(-c_c * RPARAM);
        
        # psi_r = op['HestonParams_Kappa'] * op['HestonParams_Mean'] + (1j * omega + 2.0 - type) * phi_proxy_r * op['HestonParams_Gamma'] * self.Corrs['Corr_HW2_HestVol'] * self.HullWhite2Params['Nu'] * Bf + (type - 1.0 - 1j * omega ) * phi_proxy_r * self.HestonParams['Gamma'] * self.Corrs['Corr_HW1_HestVol'] * self.HullWhite1Params['Nu'] * Bd;
        psi_r = self.HestonParams_Kappa * self.HestonParams_Mean + 1j * (omega + (type - 2)*1j) * phi_proxy_r * self.HestonParams_Gamma * self.Corr_HW2_HestVol * self.HullWhite2Params_Nu * Bf + (1.0 - 1j * (omega + (type - 2)*1j) ) * phi_proxy_r * self.HestonParams_Gamma * self.Corr_HW1_HestVol * self.HullWhite1Params_Nu * Bd;
        return func_D_r * psi_r


##    def C2Integrand(self, RPARAM, type, omega):
##        alpha   =   (-1) ** int(type)
##        Bd = (exp(-op['HullWhite1Params_Lambda'] * RPARAM)-1.0) / op['HullWhite1Params_Lambda']
##        Bf = (exp(-op['HullWhite2Params_Lambda'] * RPARAM)-1.0) / op['HullWhite2Params_Lambda']
##    
##        # Essentially D(w,r) is the same as D(w,tau), but now r is a variable 
##        a           =   0.5 * op['HestonParams_Gamma'] ** 2
##        b           =   (1j * omega + 2.0 - type) * op['Corr_Hest_HestVol'] * op['HestonParams_Gamma'] - op['HestonParams_Kappa']
##        c           =   -0.5 * (1j * alpha * omega + omega ** 2)
##        d           =   sqrt(b ** 2 - 4.0 * a * c)
##        r_plus      =   (-b + d)/(2.0 * a)
##        r_minus     =   (-b - d)/(2.0 * a)
##        g           =   r_minus / r_plus
##
##        # Now we compute Psi(r), see equation [18]. Psi(r) contains non-affine term: sqrt(v), we proxy it by phi_proxy_r, see equation [15] 
##        h = (4.0 * op['HestonParams_Kappa'] * op['HestonParams_Mean']) / op['HestonParams_Gamma'] ** 2;
##              
##        q_1 = op['HestonParams_Gamma']**2 * (1.0 - exp(-op['HestonParams_Kappa'])) / (4.0 * op['HestonParams_Kappa']);
##        l_1 = (4.0 * op['HestonParams_Kappa'] * op['HestonParams_V0'] * exp(-op['HestonParams_Kappa'])) / (op['HestonParams_Gamma']**2 * (1.0 - exp(-op['HestonParams_Kappa'])));
##        h_1 = h + l_1;
##        phi_1 = sqrt(q_1 * (l_1 - 1.0)+ q_1 * h + 0.5 * (q_1 * h) / h_1);
##        
##        a_a = sqrt(op['HestonParams_Mean'] - op['HestonParams_Gamma']**2 / (8.0 * op['HestonParams_Kappa']));
##        b_b = sqrt(op['HestonParams_V0']) - a_a;
##        c_c = -log((phi_1 - a_a) / b_b);
##        phi_proxy_r = a_a + b_b * exp(-c_c * RPARAM);
##        
##        g = (op['Corr_HW1_Hest'] * op['HullWhite1Params_Nu'] * Bd - op['Corr_HW2_Hest'] * op['HullWhite2Params_Nu'] * Bf) * phi_proxy_r + self.Corrs['Corr_HW1_HW2'] * op['HullWhite1Params_Nu'] * op['HullWhite2Params_Nu'] * Bd * Bf - 0.5 * (op['HullWhite1Params_Nu'] * Bd)**2 + (op['HullWhite2Params_Nu'] * Bf)**2
##        return g

    def Q1Calc(self, r_minus, a, g, d):
        q1 = self.HestonParams_Kappa * self.HestonParams_Mean * (r_minus * self.tenor - (1 / a) * log( (1 - g * exp(-d * self.tenor))) + (1 / a) * log((1 - g)))
        return q1
        
    def Q2Calc(self):
        h = (4 * self.HestonParams_Kappa * self.HestonParams_Mean) / self.HestonParams_Kappa**2
        q_1 = self.HestonParams_Gamma**2 * (1 - exp(-self.HestonParams_Kappa)) / (4 * self.HestonParams_Kappa)
        l_1 = (4 * self.HestonParams_Kappa * self.HestonParams_V0 * exp(-self.HestonParams_Kappa)) / (self.HestonParams_Gamma**2 * (1 - exp(-self.HestonParams_Kappa)))
        h_1 = h + l_1
        phi_1 = sqrt(q_1 * (l_1-1)+ q_1 * h + 0.5*(q_1 * h)/h_1)
    
        a_a = sqrt(self.HestonParams_Mean - self.HestonParams_Gamma** 2 / (8.0 * self.HestonParams_Kappa))
        b_b = sqrt(self.HestonParams_V0) - a_a
        c_c = -log(((phi_1)- a_a) / b_b)

        lambda_d = self.HullWhite1Params_Lambda
        lambda_f = self.HullWhite2Params_Lambda
        nu_d = self.HullWhite1Params_Nu
        nu_f = self.HullWhite2Params_Nu
        rho_ds = self.Corr_HW1_Hest
        rho_sf = self.Corr_HW2_Hest
        rho_df = self.Corr_HW1_HW2
        rho_vd = self.Corr_HW1_HestVol
        rho_vf = self.Corr_HW2_HestVol
        
        Bd_tau = (exp(-lambda_d * self.tenor)-1) / lambda_d
        Bf_tau = (exp(-lambda_f * self.tenor)-1)/lambda_f
        B2d_tau = (exp(-2 * lambda_d * self.tenor)-1) / (2 * lambda_d)
        B2f_tau = (exp(-2 * lambda_f * self.tenor)-1)/(2 * lambda_f)
        Bc_tau = (exp(-c_c * self.tenor)-1) / c_c
        Bcd_tau = (exp(-(lambda_d + c_c) * self.tenor)-1) / (lambda_d + c_c)
        Bcf_tau = (exp(-(lambda_f + c_c) * self.tenor)-1) / (lambda_f + c_c)
        Bdf_tau = (exp(-(lambda_d + lambda_f) * self.tenor)-1) / (lambda_d + lambda_f)
    
        q2a = rho_ds * nu_d * (-a_a * Bd_tau - b_b * Bcd_tau - a_a * self.tenor + b_b * Bc_tau) / lambda_d
        q2b = rho_sf * nu_f * (-a_a * Bf_tau - b_b * Bcf_tau - a_a * self.tenor + b_b * Bc_tau) / lambda_f
        q2c = rho_df * nu_d * nu_f * (self.tenor + Bd_tau + Bf_tau - Bdf_tau)/(lambda_d * lambda_f) 
        q2d = (nu_d**2) * (self.tenor + 2 * Bd_tau - B2d_tau) / (lambda_d**2) + (nu_f**2) * (self.tenor + 2 * Bf_tau - B2f_tau) / (lambda_f**2)
        q2 = q2a - q2b + q2c - 0.5 * q2d
        return q2
    
    def CharFunc(self, type, OMEGA):
        alpha   =   pow((-1), int(type))
        a           =   0.5 * self.HestonParams_Gamma ** 2
        b           =   (1j * OMEGA + 2.0 - type) * self.Corr_Hest_HestVol * self.HestonParams_Gamma - self.HestonParams_Kappa
        c           =   -0.5 * (1j * alpha * OMEGA + OMEGA ** 2)
        d           =   sqrt(b ** 2 - 4.0 * a * c)
        r_plus      =   (-b + d)/(2.0 * a)
        r_minus     =   (-b - d)/(2.0 * a)
        g           =   r_minus / r_plus
        x = log(self.initial)

        func_D = r_minus * (1.0 - exp(-d * self.tenor)) / (1 - g * exp(-d * self.tenor))

        # Then calculate C(w,tau) - separated into 2 integrals

        q2 = self.Q2Calc()        
        
        # Below combines both integrals to give C(w,tau)
        if self.isPartial: # then we can calculate func_C exactly
            q1 = self.Q1Calc(r_minus, a, g, d)
        else:
            q1 = integ(self.C1Integrand, 1e-08, self.tenor, args=(type,OMEGA))[0]
        # func_C = q1 + (1j * OMEGA + OMEGA ** 2) * q2
        func_C = q1 + (1j * (OMEGA + (type - 2)*1j) + (OMEGA + (type - 2)*1j) ** 2) * q2

        # func_E = 1j * OMEGA
        func_E = 1j * (OMEGA + (type - 2)*1j)

        return exp(func_C + func_D * self.HestonParams_V0 + func_E * x)
    
    def PIntegrand(self, OMEGA, type):
        self.nevals += 1
        x = log(self.initial)
        k = log(self.strike)
        A = self.CharFunc(type, OMEGA) * exp(-1j * OMEGA * k) * exp((type - 2) * x) / (OMEGA * 1j)
        return A.real
    
    def PFunc(self, type):
        q = integ(self.PIntegrand, 1e-08, 2000, args=(type), limit=5)[0]
        return (0.5 + (1.0/pi) * q)
    
    def getPrice(self):
        if (self.currentStatus > 0):
            op = self.allParams
            # set up all parameters upfront to avoid future dictionary lookups
            self.initial = op['initial']
            self.type = op['type']
            self.rate = op['rate']
            self.isPartial = op['isPartial']
            self.isFFT = op['isFFT']
            self.strike = op['strike']
            self.tenor = op['tenor']
            self.HullWhite1Params_Lambda = op['HullWhite1Params_Lambda']
            self.HullWhite1Params_Nu = op['HullWhite1Params_Nu']
            self.HullWhite2Params_Lambda = op['HullWhite2Params_Lambda']
            self.HullWhite2Params_Nu = op['HullWhite2Params_Nu']
            self.HestonParams_Gamma = op['HestonParams_Gamma']
            self.HestonParams_Kappa = op['HestonParams_Kappa']
            self.HestonParams_V0 = op['HestonParams_V0']
            self.HestonParams_Mean = op['HestonParams_Mean']
            self.Corr_Hest_HestVol = op['Corr_Hest_HestVol']
            self.Corr_HW1_HW2 = op['Corr_HW1_HW2']
            self.Corr_HW1_Hest = op['Corr_HW1_Hest']
            self.Corr_HW2_Hest = op['Corr_HW2_Hest']
            self.Corr_HW1_HestVol = op['Corr_HW1_HestVol']
            self.Corr_HW2_HestVol = op['Corr_HW2_HestVol']

            if op['isFFT']:
                alpha = 1.25 #arbitrary
                k = log(self.strike)
                if (self.type == 'call'):
                    val = exp(-alpha * k) * exp(-self.rate * self.tenor) / pi #* fft([FourierFunction(k, alpha)])
                else:
                    val = 0.0
            else:
                self.nevals = 0
                p1 = self.PFunc(type = 1.0)
                p2 = self.PFunc(type = 2.0)
                if (self.type == 'call'):
                    val = exp(-self.rate * self.tenor) * (self.initial * p1 - self.strike * p2)  
                else:
                    val = 0.0 # put option not done yet!
            return val